#include <math.h>
#include <stdlib.h>
#include "align.h"

#define MAX_ITER 100

/* in order to use brent minimization, we have to use several global variables here. */

static char *curr_i, *curr_j;
static DistParam *curr_dp;
static int *curr_array;
static int curr_len;

/* the matrix and the eigenvalue are copied from protdist.c in PHYLIP package */
/* this jtt matrix decomposition due to Elisabeth  Tillier */

static double jttscale = 0.01;
static double jtteigs[] =
{0.0,        -0.007031123, -0.006484345, -0.006086499, -0.005514432,
-0.00772664, -0.008643413, -0.010620756, -0.009965552, -0.011671808,
-0.012222418,-0.004589201, -0.013103714, -0.014048038, -0.003170582,
-0.00347935, -0.015311677, -0.016021194, -0.017991454, -0.018911888};

static double jttprobs[20][20] =
{{0.076999996, 0.051000003, 0.043000004, 0.051999998, 0.019999996, 0.041,
  0.061999994, 0.073999997, 0.022999999, 0.052000004, 0.090999997, 0.058999988,
  0.024000007, 0.04, 0.050999992, 0.069, 0.059000006, 0.014000008, 0.032000004,
  0.066000005},
 {0.015604455, -0.068062363, 0.020106264, 0.070723273, 0.011702977, 0.009674053,
  0.074000798, -0.169750458, 0.005560808, -0.008208636, -0.012305869,
 -0.063730179, -0.005674643, -0.02116828, 0.104586169, 0.016480839, 0.016765139,
  0.005936994, 0.006046367, -0.0082877},
 {-0.049778281, -0.007118197, 0.003801272, 0.070749616, 0.047506147,
   0.006447017, 0.090522425, -0.053620432, -0.008508175, 0.037170603,
   0.051805545, 0.015413608, 0.019939916, -0.008431976, -0.143511376,
  -0.052486072, -0.032116542, -0.000860626, -0.02535993, 0.03843545},
 {-0.028906423, 0.092952047, -0.009615343, -0.067870117, 0.031970392,
   0.048338335, -0.054396304, -0.135916654, 0.017780083, 0.000129242,
   0.031267424, 0.116333586, 0.007499746, -0.032153596, 0.033517051,
  -0.013719269, -0.00347293, -0.003291821, -0.02158326, -0.008862168},
 {0.037181176, -0.023106564, -0.004482225, -0.029899635, 0.118139633,
 -0.032298569, -0.04683198, 0.05566988, -0.012622847, 0.002023096,
 -0.043921088, -0.04792557, -0.003452711, -0.037744513, 0.020822974,
  0.036580187, 0.02331425, -0.004807711, -0.017504496, 0.01086673},
 {0.044754061, -0.002503471, 0.019452517, -0.015611487, -0.02152807,
 -0.013131425, -0.03465365, -0.047928912, 0.020608851, 0.067843095,
 -0.122130014, 0.002521499, 0.013021646, -0.082891087, -0.061590119,
  0.016270856, 0.051468938, 0.002079063, 0.081019713, 0.082927944},
 {0.058917882, 0.007320741, 0.025278141, 0.000357541, -0.002831285,
 -0.032453034, -0.010177288, -0.069447924, -0.034467324, 0.011422358,
 -0.128478324, 0.04309667, -0.015319944, 0.113302422, -0.035052393,
  0.046885372, 0.06185183, 0.00175743, -0.06224497, 0.020282093},
 {-0.014562092, 0.022522921, -0.007094389, 0.03480089, -0.000326144,
 -0.124039037, 0.020577906, -0.005056454, -0.081841576, -0.004381786,
  0.030826152, 0.091261631, 0.008878828, -0.02829487, 0.042718836,
 -0.011180886, -0.012719227, -0.000753926, 0.048062375, -0.009399129},
 {0.033789571, -0.013512235, 0.088010984, 0.017580292, -0.006608005,
 -0.037836971, -0.061344686, -0.034268357, 0.018190209, -0.068484614,
  0.120024744, -0.00319321, -0.001349477, -0.03000546, -0.073063759,
  0.081912399, 0.0635245, 0.000197, -0.002481798, -0.09108114},
 {-0.113947615, 0.019230545, 0.088819683, 0.064832765, 0.001801467,
 -0.063829682, -0.072001633, 0.018429333, 0.057465965, 0.043901014,
 -0.048050874, -0.001705918, 0.022637173, 0.017404665, 0.043877902,
 -0.017089594, -0.058489485, 0.000127498, -0.029357194, 0.025943972},
 {0.01512923, 0.023603725, 0.006681954, 0.012360216, -0.000181447,
 -0.023011838, -0.008960024, -0.008533239, 0.012569835, 0.03216118,
  0.061986403, -0.001919083, -0.1400832, -0.010669741, -0.003919454,
 -0.003707024, -0.026806029, -0.000611603, -0.001402648, 0.065312824},
 {-0.036405351, 0.020816769, 0.011408213, 0.019787053, 0.038897829,
   0.017641789, 0.020858533, -0.006067252, 0.028617353, -0.064259496,
  -0.081676567, 0.024421823, -0.028751676, 0.07095096, -0.024199434,
  -0.007513119, -0.028108766, -0.01198095, 0.111761119, -0.076198809},
 {0.060831772, 0.144097327, -0.069151377, 0.023754576, -0.003322955,
 -0.071618574, 0.03353154, -0.02795295, 0.039519769, -0.023453968,
 -0.000630308, -0.098024591, 0.017672997, 0.003813378, -0.009266499,
 -0.011192111, 0.016013873, -0.002072968, -0.010022044, -0.012526904},
 {-0.050776604, 0.092833081, 0.044069596, 0.050523021, -0.002628417,
   0.076542572, -0.06388631, -0.00854892, -0.084725311, 0.017401063,
  -0.006262541, -0.094457679, -0.002818678, -0.0044122, -0.002883973,
   0.028729685, -0.004961596, -0.001498627, 0.017994575, -0.000232779},
 {-0.01894566, -0.007760205, -0.015160993, -0.027254587, 0.009800903,
  -0.013443561, -0.032896517, -0.022734138, -0.001983861, 0.00256111,
   0.024823166, -0.021256768, 0.001980052, 0.028136263, -0.012364384,
  -0.013782446, -0.013061091, 0.111173981, 0.021702122, 0.00046654},
 {-0.009444193, -0.042106824, -0.02535015, -0.055125574, 0.006369612,
  -0.02945416, -0.069922064, -0.067221068, -0.003004999, 0.053624311,
   0.128862984, -0.057245803, 0.025550508, 0.087741073, -0.001119043,
  -0.012036202, -0.000913488, -0.034864475, 0.050124813, 0.055534723},
 {0.145782464, -0.024348311, -0.031216873, 0.106174443, 0.00202862,
  0.02653866, -0.113657267, -0.00755018, 0.000307232, -0.051241158,
  0.001310685, 0.035275877, 0.013308898, 0.002957626, -0.002925034,
 -0.065362319, -0.071844582, 0.000475894, -0.000112419, 0.034097762},
 {0.079840455, 0.018769331, 0.078685899, -0.084329807, -0.00277264,
 -0.010099754, 0.059700608, -0.019209715, -0.010442992, -0.042100476,
 -0.006020556, -0.023061786, 0.017246106, -0.001572858, -0.006703785,
  0.056301316, -0.156787357, -0.000303638, 0.001498195, 0.051363455},
 {0.049628261, 0.016475144, 0.094141653, -0.04444633, 0.005206131,
 -0.001827555, 0.02195624, 0.013066683, -0.010415582, -0.022338403,
  0.007837197, -0.023397671, -0.002507095, 0.005177694, 0.017109561,
 -0.202340113, 0.069681441, 0.000120736, 0.002201146, 0.004670849},
 {0.089153689, 0.000233354, 0.010826822, -0.004273519, 0.001440618,
  0.000436077, 0.001182351, -0.002255508, -0.000700465, 0.150589876,
 -0.003911914, -0.00050154, -0.004564983, 0.00012701, -0.001486973,
 -0.018902754, -0.054748555, 0.000217377, -0.000319302, -0.162541651}};

double *ma_transform_matrix(double mat[20][20])
{
	double *new_mat;
	int i, j, m;
	new_mat = (double*)malloc(sizeof(double) * 8000);
	for (m = 0; m < 20; ++m)
		for (i = 0; i < 20; ++i)
			for (j = 0; j < 20; ++j)
				new_mat[m + i * 20 + j * 400] = mat[m][i] * mat[m][j];
	return new_mat;
}
DistParam *ma_alloc_DistParam(int type)
{
	DistParam *dp;
	dp = (DistParam*)malloc(sizeof(DistParam));
	dp->type = type;
	dp->eigen = dp->mat = 0;
	dp->scale = 0.0;
	dp->is_kimura = 0;
	if (type == DIST_JTT) {
		dp->eigen = jtteigs;
		dp->mat = ma_transform_matrix(jttprobs);
		dp->scale = jttscale;
	} else if (type == DIST_KIMURA)
		dp->is_kimura = 1;
	return dp;
}
void ma_free_DistParam(DistParam *dp)
{
	if (dp == 0) return;
	free(dp->mat);
	free(dp);
}
void ma_cal_dist(Matrix *mat, const MultiAlign *ma, int is_rand, const DistParam *dp)
{
	extern void ma_cal_nucl(Matrix *mat, const MultiAlign *ma, int is_rand, int dist_type);
	assert(mat); assert(ma);
	if (dp->type == DIST_KIMURA || dp->type == DIST_MM || dp->type == DIST_NT_MM)
		ma_cal_mm_dist(mat, ma, is_rand, dp->is_kimura, 0);
	else if (dp->type == DIST_DN || dp->type == DIST_DS || dp->type == DIST_DM)
		ma_cal_nucl(mat, ma, is_rand, dp->type);
	else ma_cal_prob_dist(mat, ma, is_rand, dp);
}
/*
 * Calculate -log P(t).
 *
 * This function is greatly optimized because 98% time will be spent on this
 * function. First, curr_dp->mat is pre-calculated, this will avoid one float
 * times. Second, see the prob calculation that save a log().
 */
static double cal_prob(double t) /* calculate likelihood given distance */
{
	int l, m, c1, c2, len, count;
	double *tmp_ptr;
	double tmp[20], p, prob, log_prob;

	for (m = 0; m < 20; ++m) tmp[m] = exp(curr_dp->eigen[m] * t);
	prob = 1.0; count = 0;
	len = curr_len;
	for (l = 0; l < len; ++l) {
		c1 = curr_i[curr_array[l]]; c2 = curr_j[curr_array[l]];
		if (c1 >= 20 || c2 >= 20) continue;
		p = 0.0;
		tmp_ptr = curr_dp->mat + c1 * 20 + c2 * 400;
		for (m = 0; m < 20; ++m)
			p += tmp_ptr[m] * tmp[m];
		prob *= p;
		if (prob < 1e-100) {
			prob /= 1e-100;
			++count;
		}
	}
	log_prob = -(log10(prob) - 100 * count);
	return log_prob;
}
void ma_cal_prob_dist(Matrix *mat, const MultiAlign *ma, int is_rand, const DistParam *dp)
{
	extern double brent(double ax, double bx, double cx, double (*f)(double), double tol, double *xmin);
	int i, j, l, n;
	int len, *array;
	double t, max_t;

	max_t = -10.0;
	n = ma->n;
	len = ma->len;

	array = (int*)malloc(sizeof(int) * len);
	if (is_rand) {
		for (l = 0; l < len; ++l)
#ifdef _WIN32
			array[l] = (int)((double)rand() / RAND_MAX * len);
#else
			array[l] = (int)(drand48() * len);
#endif
	} else for (l = 0; l < len; ++l) array[l] = l;
	ma_cal_mm_dist(mat, ma, is_rand, 1, array); /* calculate Kimura distance */
	curr_dp = (DistParam*)dp;
	curr_array = array;
	curr_len = len;

	for (i = 0; i < n; ++i) {
		curr_i = ma->seq[i];
		for (j = 0; j < i; ++j) {
			curr_j = ma->seq[j];
			if (mat->dist[i * n + j] < 0.0) continue; /* fill it later */
			brent(MA_MIN_DIST / dp->scale, mat->dist[i * n + j] / dp->scale, MA_MAX_DIST / dp->scale,
				  cal_prob, MA_MIN_DIST / 10.0 / dp->scale, &t);
			mat->dist[i * n + j] = mat->dist[j * n + i] = t * dp->scale;
			if (max_t < t) max_t = t;
		}
	}
	max_t *= 2.0 * dp->scale;
	for (i = 0; i < n * n; ++i) /* handle the undefined distances */
		if (mat->dist[i] < 0.0) mat->dist[i] = max_t;
	for (i = 0; i < n; ++i)
		mat->dist[i * n + i] = 0.0;
	free(array);
}
